import java.util.ArrayList;
import java.util.Comparator;
import java.util.List;
import java.util.PriorityQueue;

public class KSmallestPairs {

    public List<List<Integer>> kSmallestPairs1(int[] nums1, int[] nums2, int k) {
        PriorityQueue<List<Integer>> heap = new PriorityQueue<>(new Comparator<List<Integer>>() {
            @Override
            public int compare(List<Integer> o1, List<Integer> o2) {
                return (o2.get(0) + o2.get(1)) - (o1.get(0) + o1.get(1));
            }
        });
        for (int i = 0; i < nums1.length; i++) {
            for (int j = 0; j < nums2.length; j++) {
                List<Integer> list = heap.peek();
                if (heap.size() >= k && list.get(0) + list.get(1) <= nums1[i] + nums2[j]) continue;
                List<Integer> temp = new ArrayList<>();
                temp.add(nums1[i]);
                temp.add(nums2[j]);
                if (heap.size() >= k) heap.poll();
                heap.add(temp);
            }
        }
        List<List<Integer>> ans = new ArrayList<>();
        for (int i = 0; i < k; i++) {
            if (!heap.isEmpty()) ans.add(heap.poll());
        }
        return ans;
    }

    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        PriorityQueue<int[]> heap = new PriorityQueue<>(new Comparator<int[]>() {
            @Override
            public int compare(int[] o1, int[] o2) {
                return nums1[o1[0]] + nums2[o1[1]] - nums1[o2[0]] - nums2[o2[1]];
            }
        });
        for (int i = 0; i < Math.min(k, nums1.length); i++) {
            heap.add(new int[]{i, 0});
        }
        List<List<Integer>> ans = new ArrayList<>();
        while (!heap.isEmpty() && k-- > 0) {
            int[] poll = heap.poll();
            ArrayList<Integer> list = new ArrayList<>();
            list.add(nums1[poll[0]]);
            list.add(nums2[poll[1]]);
            ans.add(list);
            if (poll[1] + 1 < nums2.length) heap.add(new int[]{poll[0], poll[1] + 1});
        }
        return ans;
    }
}